# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: chensili

import time
import logging
import json
import os

from absl import app, flags
import hvac

import coin.base.database.util as db_util
import coin.support.accounting.logic.account_status_util as account_status
from coin.base.param_util import split_string
from coin.support.accounting.logic.query_util import (
    filter_invalid_request_and_result)
from coin.support.accounting.app.account_info_collector import (
    _load_proxy_config,
    generate_account_loggers)
from coin.support.accounting.logic.key_util import get_key_file
from xunkemgmt_client.support.accounting.database.mysql_importer import \
    MysqlAccountInfoImporter as MysqlImporter
from xunkemgmt_client.support.accounting.database.memcached_importer import (
    MemcachedAccountInfoImporter)
from xunkemgmt_client.client.util.query_util import query_accounts


FLAGS = flags.FLAGS


def _dump_coin_key_from_vault(request_list, vault_addr, vault_token_file):
  key_files = [
      get_key_file(request.market_type, request.exchange, request.owner,
                    request.api_version, "view", check_exist=False)
      for request in request_list
  ]
  with open(FLAGS.vault_token_file) as fp:
    vault_token = json.load(fp)['token']
  client = hvac.Client(url=FLAGS.vault_addr, token=vault_token)
  assert client.is_authenticated()
  for key_file in key_files:
    key_path = os.path.relpath(key_file, os.path.expanduser(FLAGS.key_root))
    key_content = client.secrets.kv.read_secret_version(
        path=os.path.join('view', key_path), mount_point='coin_key')['data']['data']
    out_path = os.path.join(
        os.path.abspath(os.path.expanduser(FLAGS.key_root)), key_path)
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    with open(out_path, 'w') as fp:
      json.dump(key_content, fp, indent=2)


def collect_balances(request_list, proxy, account_loggers):
  if len(request_list) == 0:
    return

  logging.info('start balance query')
  balance_response_list = account_status.query_balance_list(request_list, proxy)
  balance_request_list, balance_response_list = \
    filter_invalid_request_and_result(request_list, balance_response_list)

  logging.info('start account log writing')
  for request, response in zip(balance_request_list, balance_response_list):
    account_logger = account_loggers[
        (request.market_type, request.exchange, request.owner)]
    account_logger.write_balance(response.balance)

  # logging.info('start balance insertion')
  # db_config = db_util.read_db_config(FLAGS.memcached_config)
  # importer = MemcachedAccountInfoImporter(db_config=db_config)
  # importer.insert_account_balance_history(balance_request_list, balance_response_list)

  # db_config = db_util.read_db_config(FLAGS.mysql_config)
  # importer = MysqlImporter(db_config=db_config)
  # importer.insert_account_balance_history(balance_request_list, balance_response_list)


def collect_positions(request_list, proxy, account_loggers):
  request_list = [req for req in request_list if req.market_type == 'Futures']
  if len(request_list) == 0:
    return

  logging.info('start position query')
  position_response_list = account_status.query_position_list(position_request_list, proxy)
  position_request_list, position_response_list = \
    filter_invalid_request_and_result(position_request_list, position_response_list)

  logging.info('start account log writing')
  for request, response in zip(position_request_list, position_response_list):
    account_logger = account_loggers[
        (request.market_type, request.exchange, request.owner)]
    account_logger.write_position(response.position)

  # logging.info('start position insertion')
  # db_config = db_util.read_db_config(FLAGS.memcached_config)
  # importer = MemcachedAccountInfoImporter(db_config=db_config)
  # importer.insert_account_position_history(position_request_list, position_response_list)

  # db_config = db_util.read_db_config(FLAGS.mysql_config)
  # importer = MysqlImporter(db_config=db_config)
  # importer.insert_account_position_history(position_request_list, position_response_list)


def main(_):
  start_time = time.time()
  request_list = query_accounts(
      active=True,
      market_types=split_string(FLAGS.market_type),
      exchanges=split_string(FLAGS.exchange),
      owners=split_string(FLAGS.owner),
      as_proto=True)
  assert len(request_list) > 0
  proxy = _load_proxy_config(FLAGS.proxy_config_file)
  if FLAGS.vault_addr is not None:
    _dump_coin_key_from_vault(request_list, FLAGS.vault_addr, FLAGS.vault_token_file)
  account_loggers = generate_account_loggers(
      request_list, log_root=os.path.expanduser('~/data/account_proto_log'))

  while True:
    collect_balances(request_list, proxy, account_loggers)
    collect_positions(request_list, proxy, account_loggers)
    logging.info("end balance collector")
    if FLAGS.duration is not None and \
       time.time() - start_time > FLAGS.duration:
      break
    time.sleep(FLAGS.query_frequency)


if __name__ == '__main__':
  logging.basicConfig(level='INFO', format='%(levelname)8s %(asctime)s %(message)s')

  flags.DEFINE_string('market_type', None, 'Comma separated market_type.')

  flags.DEFINE_string('exchange', None, 'Comma separated exchange.')

  flags.DEFINE_string('owner', None, 'Comma separated account owner.')

  flags.DEFINE_string('api_version', None, 'Comma separated api version.')

  flags.DEFINE_integer('query_frequency', 30, 'query frequency in seconds.')

  flags.DEFINE_integer(
      'duration',
      None,
      'Duration for running collector in seconds. Eg. If set to 300, program will '
      'stop after running 300 seconds. If set to None, program will run forever')

  flags.DEFINE_string('proxy_config_file', None, 'proxy_config_file')

  flags.DEFINE_string('mysql_config',
                      "../../coin_deploy/support_accounting/db_config/mysql_config_prod.json",
                      'path of database config.')

  flags.DEFINE_string('memcached_config',
                      "../../coin_deploy/support_accounting/db_config/memcached_config_prod.json",
                      'path of database config.')

  flags.DEFINE_string('vault_addr', None, '')
  flags.DEFINE_string('vault_token_file', '../../coin_key/vault/coin_key_high.json', '')

  app.run(main)
